#!/usr/bin/perl
#
# conntrack_viewer.pl
#
# Copyright (C) 2009 Daniel Meekins
# Contact: dmeekins - gmail
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# See the README file for information about this script.
#

use strict;
use Switch;
use Curses;

############################################################################
# Function: usage
# Args: none
# Summary: Prints arguments required for the script.
#
sub usage
{
    print "usage: ${0} [ip_conntrack file]\n";
    print "If no file is specified, will default to /proc/net/ip_conntrack\n";
    exit;
}

############################################################################
# Function: print_screen
# Args:
#     $start_num - Line number to start at for printing all the lines in the
#                  passed array.
#     @arr - Array containing all the lines to print, already formated.
# Returns: none
# Summary: Print the most lines that can fit on the window, including the
#          command information.
#
sub print_screen
{
    my($start_num, @arr) = @_;

    my $count = 0;
    my $max_count = getmaxy($stdscr) - 4; # 4 less lines b/c of header and
                                          # list of commands.

    clear();
    move(0, 0);

    # Print the table header
    printw("Source          Port  Destination     Port  ");
    printw("Proto State       TTL\n");
    
    # Print each line, skipping $start_num lines before printing the first
    # one, up to $max_count number of lines.
    foreach(@arr)
    {
        if(($count >= $start_num) && ($count - $start_num < $max_count))
        {
            printw($_);
        }
        
        $count += 1;
    }

    # Print the command info at the bottom.
    move(getmaxy($stdscr) - 3, 0);
    printw("Sort by: src (1), dest (2), sport (3), dport (4), ");
    printw("proto (5), state (6)\n");
    printw("Filter by: src (a), dest (b), sport (c), dport (d), ");
    printw("proto (e), state (f)\n");
    printw("Undo: src (A), dest (B), sport (C), ");
    printw("dport (D), proto (E), state (F), all (G)");

    refresh();
}

############################################################################
# Function: get_line
# Args: 
#     $query - Line to print at bottom of screen
# Return: String entered by the user.
# Summary: Print the $query line at the bottom of the screen and waits for
#          the user to supply a string of input, finished by a newline.
#
sub get_line
{
    my($query) = @_;
    my $str = "";

    timeout(-1); # enable blocking input

    # Print the query line, with a ": " after it.
    move(getmaxy($stdscr)-1, 0);
    printw(" "x(getmaxx($stdscr)-1));
    move(getmaxy($stdscr)-1, 0);
    printw($query.": ");
    refresh();

    getstr($str);

    timeout(100); # reset input to nonblocking

    return $str;
}

############################################################################
# Function: by_ip_addr
# Args: 
#     $a, $b - IP address string, in the format of #.#.#.#.# (last # was
#              appended in loop in main() to differentiate between repeating
#              addresses.
# Return: 0 if $a and $b are the same, a negative integer if $a is less than
#         $b, or a positive integer if $ is greater than $b.
# Summary: Compares the two IP addresses by numbers. Used as the subroutine
#          supplied to the sort() command.
#
sub by_ip_addr
{
    my @a_arr = split(/\./, $a);
    my @b_arr = split(/\./, $b);
    my @a_nums = ();
    my @b_nums = ();
    my $s;

    # Populate the integer array from the string array, converting the strings
    # to ints.
    foreach $s (@a_arr)
    {
        push(@a_nums, int($s));
    }

    foreach $s (@b_arr)
    {
        push(@b_nums, int($s));
    }

    # Compare each number from the ints parsed from the dot-decimal notation,
    # from left to right. If both IP's were the same, then compare the number
    # appended at the end.
    if($a_nums[0] == $b_nums[0])
    {
        if($a_nums[1] == $b_nums[1])
        {
            if($a_nums[2] == $b_nums[2])
            {
                if($a_nums[3] == $b_nums[3])
                {
                    return $a_nums[4] - $b_nums[4];
                }
                else
                {
                    return $a_nums[3] - $b_nums[3];
                }
            }
            else
            {
                return $a_nums[2] - $b_nums[2];
            }
        }
        else
        {
            return $a_nums[1] - $b_nums[1];
        }
    }
    else
    {
        return $a_nums[0] - $b_nums[0];
    }
}

############################################################################
# Function: by_port
# Args:
#     $a, $b, - String representing the port number, with a ".#" appened to
#               the end in case there were repeating port numbers.
# Return: 0 if the numbers are the same, a negative integer if $a is less than
#         $b, or a positive integer if $a is greater than $b.
# Summary: Compares two port numbers numerically. Used as the sorting routine
#          supplied to the sort() command.
#
sub by_port
{
    my $porta;
    my $portb;
    my $numa;
    my $numb;

    # Parse the port string to the port number and repeating number.
    ($porta, $numa) = split(/\./, $a);
    ($portb, $numb) = split(/\./, $b);

    if($porta eq $portb)
    {
        return int($numa) - int($numb);
    }
    else
    {
        return int($porta) - int($portb);
    }
}

############################################################################
# Function: main
# Args: (uses @ARGV)
# Return: none
# Summary: Main control routine that continously reads the conntrack file,
#          parses and displays it, and handles input from the user.
#
sub main
{
    my $filter_src = "";
    my $filter_dst = "";
    my $filter_sport = "";
    my $filter_dport = "";
    my $filter_proto = "";
    my $filter_state = "";
    my $sortby = "src";

    my $filename = $ARGV[0];
    my $running = 1;
    my $line_start = 0;

    if($ARGV[0] eq "-h")
    {
        usage();
        return;
    }
    
    if($filename eq "")
    {
        $filename = "/proc/net/ip_conntrack"
    }
    
    initscr();
    cbreak();
    keypad($stdscr, 1); # Enable to can use arrow and PGUP/PGDN keys.
    timeout(100);       # Set non-blocking timeout to 100ms.
    
    while($running)
    {
        my @lines;
        my %str_hash = ();
        my @sorted_arr = ();
        my $max_lines;
        my $num_lines;
        my $c;

        # used for loops
        my $count = 0;
        my $line;
        my $str;
        my $key;

        # Immediately get all the file contents and close it.
        open(CONN_FILE, "<${filename}") || goto done;
        @lines = <CONN_FILE>;
        close(CONN_FILE);

        # Parse all the connection lines and add then to the hash table.
        foreach $line (@lines)
        {
            # NOTE: To use the regex with the ip_conntrack file, comment the
            # lines with [nf_conntrack] in the comment. To use with the
            # nf_conntrack file, uncomment those lines and comment the
            # line with [ip_conntrack] in the comment.
            if($line =~ m/^
                           #\w+              # IP version      [nf_conntrack]
                           #\ {1,8}\d+       # IP ver # (?)    [nf_conntrack]
                           #\ (\w+)          # Protocol ($1)   [nf_conntrack]
                           (\w+)            # Protocol ($1)   [ip_conntrack]
                           \ {1,8}\d+       # 
                           \ (\d+)          # TTL ($2)
                           \ (\w*)          # State ($3)
                           \ ?src=([\d\.]+) # Src addr ($4)
                           \ dst=([\d\.]+)  # Dst addr ($5)
                           \ sport=(\d+)    # Src port ($6)
                           \ dport=(\d+)    # Dst port ($7)
                           \ .*$/x)         # ... ignore rest of string
            {
                # Add the line to the hash table only if it passes the filter.
                if(($filter_src eq "" || $filter_src eq $4) &&
                   ($filter_dst eq "" || $filter_dst eq $5) &&
                   ($filter_sport eq "" || $filter_sport eq $6) &&
                   ($filter_dport eq "" || $filter_dport eq $7) &&
                   ($filter_proto eq "" || $filter_proto eq $1) &&
                   ($filter_state eq "" || $filter_state eq $3))
                {
                    # Format to match the table header (see print_screen
                    # function).
                    $str = sprintf("%-15s %-5s %-15s %-5s %-5s %-11s %-5s\n", 
                                   $4, $6, $5, $7, $1, $3, $2);

                    # Use one of the fields as the key to the hash table,
                    # according to how it should be sorted.
                    switch($sortby)
                    {
                        case "src" { $str_hash{$4.".$count"} = $str; }
                        case "dst" { $str_hash{$5.".$count"} = $str; }
                        case "sport" { $str_hash{$6.".$count"} = $str; }
                        case "dport" { $str_hash{$7.".$count"} = $str; }
                        case "proto" { $str_hash{$1.".$count"} = $str; }
                        case "state" { $str_hash{$3.".$count"} = $str; }
                    }
                }
                
                $count += 1;
            }
        }

        # Sort the hash table, putting the sorted values into the @sorted_arr
        # array. Use one of the specific sorting functions if need to sort
        # by IP address or port.
        if($sortby eq "src" || $sortby eq "dst")
        {
            foreach $key (sort by_ip_addr keys %str_hash)
            {
                push(@sorted_arr, $str_hash{$key});
            }
        }
        elsif($sortby eq "sport" || $sortby eq "dport")
        {
            foreach $key (sort by_port keys %str_hash)
            {
                push(@sorted_arr, $str_hash{$key});
            }
        }
        else
        {
            foreach $key (sort keys %str_hash)
            {
                push(@sorted_arr, $str_hash{$key});
            }
        }
        
        $max_lines = getmaxy($stdscr) - 4;
        $num_lines = scalar @sorted_arr;

        # Get any user input waiting
        $c = getch();

        # Check if one of the scrolling keys were pressed, to adjsut the view
        # before printing to the screen.
        switch($c)
        {
            case KEY_DOWN { $line_start += 1; }
            case KEY_UP { $line_start -= 1; }
            case KEY_NPAGE { $line_start += $max_lines; }
            case KEY_PPAGE { $line_start -= $max_lines; }
        }
        
        # Adjust $line_start if it went under or over the number of connection
        # lines.
        if($line_start < 0)
        {
            $line_start = 0;
        }
        elsif($line_start >= $num_lines)
        {
            $line_start = $num_lines - $max_lines - 1;
        }
        
        print_screen($line_start, @sorted_arr);
        sleep(0.7);

        # If the user typed in some command key, then handle that to use for
        # the next time through the loop. See print_screen() for the info
        # that is printed about the commands.
        switch($c)
        {
            case "q" { $running = 0; }
            case "a" { 
                $filter_src = get_line("source IP address");
                if($filter_src !~ /^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}$/)
                {
                    $filter_src = "";
                }
            }
            case "b" { 
                $filter_dst = get_line("destination IP address");
                if($filter_dst !~ /^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}$/)
                {
                    $filter_dst = "";
                }
            }
            case "c" { 
                $filter_sport = get_line("source port");
                if($filter_sport !~ /^\d{1,5}$/)
                {
                    $filter_sport = "";
                }
            }
            case "d" { 
                $filter_dport = get_line("destination port");
                if($filter_dport !~ /^\d{1,5}$/)
                {
                    $filter_dport = "";
                }
            }
            case "e" { 
                $filter_proto = get_line("protocol");
                if($filter_proto ne "tcp" && $filter_proto ne "udp" &&
                   $filter_proto ne "icmp")
                {
                    $filter_proto = "";
                }
            }
            case "f" { $filter_state = get_line("connection state"); }
            case "A" { $filter_src = ""; }
            case "B" { $filter_dst = ""; }
            case "C" { $filter_sport = ""; }
            case "D" { $filter_dport = ""; }
            case "E" { $filter_proto = ""; }
            case "F" { $filter_state = ""; }
            case "G" {
                $filter_src = "";
                $filter_dst = "";
                $filter_sport = "";
                $filter_dport = "";
                $filter_proto = "";
                $filter_state = "";
            }
            case "1" { $sortby = "src"; }
            case "2" { $sortby = "dst"; }
            case "3" { $sortby = "sport"; }
            case "4" { $sortby = "dport"; }
            case "5" { $sortby = "proto"; }
            case "6" { $sortby = "state"; }
        } 
    }
    
  done:
    endwin(); # must do this.
}

main();
